import onnx
import onnx_graphsurgeon as gs
import numpy as np

def add_topk_postprocess(model_path, num_classes, num_keypoints, k, output_model_path, had_visible):
    model = onnx.load(model_path)
    graph = gs.import_onnx(model)

    # 获取输出节点
    output_node = graph.outputs[0]
    output_shape = output_node.shape
    batch_size, channels, num_boxes = output_shape
    print("原输出层形状:", output_shape)

    # Transpose 操作： (batch_size, 25, num_boxes) -> (batch_size, num_boxes, 25)
    transposed_output = gs.Variable(name="transposed_output", dtype=np.float32, shape=[batch_size, num_boxes, channels])
    transpose_node = gs.Node(op="Transpose",
                             inputs=[output_node],
                             outputs=[transposed_output],
                             attrs={"perm": [0, 2, 1]})  # 交换第2、3维度
    graph.nodes.append(transpose_node)
    print(transpose_node.outputs)

    # 分解输出
    # boxes: 4 (center_x, center_y, width, height)
    # confidences: num_classes (class1_conf, class2_conf, class3_conf)
    # keypoints: num_keypoints * 3 (x, y, vis) * 6
    boxes = gs.Variable(name="boxes", dtype=np.float32, shape=[batch_size, num_boxes, 4])
    confidences = gs.Variable(name="confidences", dtype=np.float32, shape=[batch_size, num_boxes, num_classes])
    if had_visible:
        keypoints = gs.Variable(name="keypoints", dtype=np.float32, shape=[batch_size, num_boxes, num_keypoints * 3])
    else:
        keypoints = gs.Variable(name="keypoints", dtype=np.float32, shape=[batch_size, num_boxes, num_keypoints * 2])

    # Split 操作：沿最后一个维度进行分割
    if had_visible:
        split_node = gs.Node(op="Split",
                             inputs=[transposed_output],
                             outputs=[boxes, confidences, keypoints],
                             attrs={"axis": -1, "split": [4, num_classes, num_keypoints * 3]})
    else:
        split_node = gs.Node(op="Split",
                             inputs=[transposed_output],
                             outputs=[boxes, confidences, keypoints],
                             attrs={"axis": -1, "split": [4, num_classes, num_keypoints * 2]})
    graph.nodes.append(split_node)
    print(split_node.outputs)

    # 获取最大置信度类别和最大置信度
    max_scores = gs.Variable(name="max_scores", dtype=np.float32, shape=[batch_size, num_boxes, 1])
    max_classes = gs.Variable(name="max_classes", dtype=np.int64, shape=[batch_size, num_boxes, 1])
    argmax_node = gs.Node(op="ArgMax",
                          inputs=[confidences],
                          outputs=[max_classes],
                          attrs={"axis": -1, "keepdims": 1})
    max_node = gs.Node(op="ReduceMax",
                       inputs=[confidences],
                       outputs=[max_scores],
                       attrs={"axes": [-1], "keepdims": 1})
    graph.nodes.extend([argmax_node, max_node])
    print(argmax_node.outputs)
    print(max_node.outputs)

    # 5. Shuffle层进行维度变换

    shuffle_scores = gs.Variable(name="shuffle_scores", dtype=np.float32, shape=[batch_size, num_boxes])
    shuffle_shape = gs.Constant(name="shuffle_shape", values=np.array([batch_size, num_boxes], dtype=np.int64))
    shuffle_node = gs.Node(op='Reshape',
                           inputs=[max_scores, shuffle_shape],
                           outputs=[shuffle_scores])
    graph.nodes.append(shuffle_node)
    print(shuffle_node.outputs)

    # 6. TopK层进行Top-K操作
    k_top = gs.Constant(name="k_top", values=np.array([k], dtype=np.int64))
    topk_values = gs.Variable(name="topk_values", dtype=np.float32, shape=[batch_size, k])
    topk_indices = gs.Variable(name="topk_indices", dtype=np.int64, shape=[batch_size, k])
    topk_node = gs.Node(op='TopK',
                        inputs=[shuffle_scores, k_top],
                        outputs=[topk_values, topk_indices],
                        attrs={"largest": 1, "sorted": 1})
    graph.nodes.append(topk_node)
    print(topk_node.outputs)

    # 7. 使用Gather层进行数据聚合


    max_classes_float = gs.Variable(name="final_classes_float", dtype=np.float32, shape=[batch_size, num_boxes, 1])
    cast_classes = gs.Node(op="Cast",
                           inputs=[max_classes],
                           outputs=[max_classes_float],
                           attrs={"to": onnx.TensorProto.FLOAT})
    graph.nodes.append(cast_classes)

    if had_visible:
        concat_output = gs.Variable(name="concat_output", dtype=np.float32, shape=[batch_size, num_boxes, 4 + 1 + 1 + num_keypoints * 3])
    else:
        concat_output = gs.Variable(name="concat_output", dtype=np.float32, shape=[batch_size, num_boxes, 4 + 1 + 1 + num_keypoints * 2])
    concat_node = gs.Node(op="Concat",
                          inputs=[boxes, max_scores, max_classes_float, keypoints],
                          outputs=[concat_output],
                          attrs={"axis": -1})
    graph.nodes.append(concat_node)
    print(concat_node.outputs)


    if had_visible:
        gather_output = gs.Variable(name="gather_output", dtype=np.float32, shape=[batch_size, 1, k, 4 + 1 + 1 + num_keypoints * 3])
    else:
        gather_output = gs.Variable(name="gather_output", dtype=np.float32, shape=[batch_size, 1, k, 4 + 1 + 1 + num_keypoints * 2])
    gather_node = gs.Node(op='Gather',
                          inputs=[concat_output, topk_indices],
                          outputs=[gather_output],
                          attrs={"axis": 1})
    graph.nodes.append(gather_node)
    print(gather_node.outputs)


    if had_visible:
        final_output = gs.Variable(name="final_output", dtype=np.float32, shape=[batch_size, k, 4 + 1 + 1 + num_keypoints * 3])
        shuffle_shape2 = gs.Constant(name="shuffle_shape2", values=np.array([batch_size, k, 4 + 1 + 1 + num_keypoints * 3], dtype=np.int64))
    else:
        final_output = gs.Variable(name="final_output", dtype=np.float32, shape=[batch_size, k, 4 + 1 + 1 + num_keypoints * 2])
        shuffle_shape2 = gs.Constant(name="shuffle_shape2", values=np.array([batch_size, k, 4 + 1 + 1 + num_keypoints * 2], dtype=np.int64))

    shuffle_node2 = gs.Node(op='Reshape',
                           inputs=[gather_output, shuffle_shape2],
                           outputs=[final_output])
    graph.nodes.append(shuffle_node2)
    print(shuffle_node2.outputs)

    # 更新输出
    graph.outputs = [final_output]


    # 整理图结构并保存
    graph.cleanup().toposort()
    onnx.save(gs.export_onnx(graph), output_model_path)
    print("成功保存带 NMS 和 Keypoints 合并输出的 ONNX 模型")

